﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Graphics;
using Microsoft.Xna.Framework.Content.Pipeline;
using Microsoft.Xna.Framework.Content.Pipeline.Graphics;
using Microsoft.Xna.Framework.Content.Pipeline.Processors;

namespace SkinnedModelProcessor
{
    [ContentProcessor]
    public class SkinnedModelProcessor : ModelProcessor
    {
        public override ModelContent Process(NodeContent input, ContentProcessorContext context)
        {
            ValidateMesh(input, context, null);

            //find the skeleton
            BoneContent skeleton = MeshHelper.FindSkeleton(input);
            if (skeleton == null)
            {
                throw new InvalidContentException("Input skeleton not found.");
            }
            FlattenTransforms(input, skeleton);
            //read the bind pose and skeleton hierarchy data
            IList<BoneContent> bones = MeshHelper.FlattenSkeleton(skeleton);
            if (bones.Count > SkinnedEffect.MaxBones)
            {
                throw new InvalidContentException(string.Format("Skeleton has {0} bones, but the maximum supported is {1}.", bones.Count, SkinnedEffect.MaxBones));
            }
            List<Matrix> bindPose = new List<Matrix>();
            List<Matrix> inverseBindPose = new List<Matrix>();
            List<int> skeletonHierarchy = new List<int>();
            foreach (BoneContent bone in bones)
            {
                bindPose.Add(bone.Transform);
                inverseBindPose.Add(Matrix.Invert(bone.AbsoluteTransform));
                skeletonHierarchy.Add(bones.IndexOf(bone.Parent as BoneContent));
            }
            //convert animation data to our runtime format
            Dictionary<string, AnimationClip> animationClips;
            animationClips = ProcessAnimations(skeleton.Animations, bones);

            //chain to the base ModelProcessor class
            ModelContent model = base.Process(input, context);
            model.Tag = new SkinningData(animationClips, bindPose, inverseBindPose, skeletonHierarchy);
            return model;
        }
        static void FlattenTransforms(NodeContent node, BoneContent skeleton)
        {
            foreach (NodeContent child in node.Children)
            {
                //don't process skeleton
                if (child == skeleton) continue;
                //bake the local transform inte the actual geometry
                MeshHelper.TransformScene(child, child.Transform);
                //set local coordinates back to identity
                child.Transform = Matrix.Identity;
                //recurse
                FlattenTransforms(child, skeleton);
            }
        }
        static Dictionary<string, AnimationClip> ProcessAnimations(AnimationContentDictionary animations, IList<BoneContent> bones)
        {
            //build a table mapping bone names to indices
            Dictionary<string, int> boneMap = new Dictionary<string, int>();
            for (int i = 0; i < bones.Count; i++)
            {
                string boneName = bones[i].Name;
                if (!string.IsNullOrEmpty(boneName))
                    boneMap.Add(boneName, i);
            }
            //convert each animation in turn
            Dictionary<string, AnimationClip> animationClips;
            animationClips = new Dictionary<string, AnimationClip>();
            foreach (KeyValuePair<string, AnimationContent> animation in animations)
            {
                AnimationClip processed = ProcessAnimation(animation.Value, boneMap);
                animationClips.Add(animation.Key, processed);
            }
            if (animationClips.Count == 0)
            {
                throw new InvalidContentException("Input file does not contain any animations.");
            }
            return animationClips;
        }
        static AnimationClip ProcessAnimation(AnimationContent animation, Dictionary<string, int> boneMap)
        {
            List<Keyframe> keyframes = new List<Keyframe>();
            //for each input animation channel
            foreach (KeyValuePair<string, AnimationChannel> channel in animation.Channels)
            {
                //look up what bone this channel is controlling
                int boneIndex;
                if (!boneMap.TryGetValue(channel.Key, out boneIndex))
                {
                    throw new InvalidContentException(string.Format("Found animation for bone '{0}', " +
                        "which is not part of the skeleton.", channel.Key));
                }
                //convert the keyframe data
                foreach (AnimationKeyframe keyframe in channel.Value)
                {
                    keyframes.Add(new Keyframe(boneIndex, keyframe.Time, keyframe.Transform));
                }
            }
            //sort the merget keyframes by time
            keyframes.Sort(CompareKeyframeTimes);
            if (keyframes.Count == 0)
            {
                throw new InvalidContentException("Animation has no keyframes.");
            }
            if (animation.Duration <= TimeSpan.Zero)
            {
                throw new InvalidContentException("Animation has a zero duration.");
            }
            return new AnimationClip(animation.Duration, keyframes);
        }
        static int CompareKeyframeTimes(Keyframe a, Keyframe b)
        {
            return a.Time.CompareTo(b.Time);
        }
        static void ValidateMesh(NodeContent node, ContentProcessorContext context, string parentBoneName)
        {
            MeshContent mesh = node as MeshContent;
            if (mesh != null)
            {
                if (parentBoneName != null)
                {
                    context.Logger.LogWarning(null, null, "Mesh {0} is a child of bone {1}. SkinnedModelProcessor " +
                        "doesn't correctly handle meshes that are children of bones.", mesh.Name, parentBoneName);
                }
                if (!MeshHasSkinning(mesh))
                {
                    context.Logger.LogWarning(null, null,
                        "Mesh {0} has no skinning info, so it has been deleted.",
                        mesh.Name);
                    mesh.Parent.Children.Remove(mesh);
                    return;
                }
            }
            else if (node is BoneContent)
            {
                //if this is a bone, remember that we are now looking inside it.
                parentBoneName = node.Name;
            }
            //recurse (iterating over a copy of the child collection, because
            //validating children may delete some of them).
            foreach (NodeContent child in new List<NodeContent>(node.Children))
            {
                ValidateMesh(child, context, parentBoneName);
            }
        }
        static bool MeshHasSkinning(MeshContent mesh)
        {
            foreach (GeometryContent geometry in mesh.Geometry)
            {
                if (!geometry.Vertices.Channels.Contains(VertexChannelNames.Weights()))
                {
                    return false;
                }
            }
            return true;
        }
    }
}
